/**
 * 
 */
package com.gragra.sampling.nonproj;
import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntIterator;
import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
import com.gragra.sampling.LogarithmicFactorialPlusAlphaCash;
import com.gragra.sampling.PopulationSamplingStructure;
import com.gragra.sampling.binary.HeadAssignmentStructure;
import com.gragra.sampling.nonproj.DMVRoot;
import com.gragra.sampling.nonproj.structures.nonprojective.NonProjDepdTree;
import com.gragra.sampling.sequentialStructure.WordAndTagHolder;
import com.gragra.util.OpenUnitElementCounter;
import com.gragra.util.StringToInt;
/**
 * @author Christoph Teichmann
 * 23.10.2013
 */
public class TwoLevelDMV extends DMVRoot
{

	public TwoLevelDMV(int population, IntSet tags, double headHeadAlpha,
			double stopAlpha, double functionWordConstraint,
			double nounWordConstraint, double structMax, double structRed)
	{
		super(population, tags, headHeadAlpha, stopAlpha, functionWordConstraint,
				nounWordConstraint, structMax, structRed);
		this.sSecUnit = new LogarithmicFactorialPlusAlphaCash(stopAlpha*2, population);
		this.hSecUnit = new LogarithmicFactorialPlusAlphaCash(headHeadAlpha*tags.size()/2, population);
		this.hSecElement = new LogarithmicFactorialPlusAlphaCash(headHeadAlpha, population);
		this.sSecElement = new LogarithmicFactorialPlusAlphaCash(stopAlpha, population);
		IntIterator iit = tags.iterator();
		this.headToHeadLeftSecUnits = new Int2IntOpenHashMap();
		this.headToHeadRightSecUnits = new Int2IntOpenHashMap();
		this.stopAdjacentLeftSecUnits = new Int2IntOpenHashMap();
		this.stopAdjacentRightSecUnits = new Int2IntOpenHashMap();
		this.stopDistantLeftSecUnits = new Int2IntOpenHashMap();
		this.stopDistantRightSecUnits = new Int2IntOpenHashMap();
		IntSet seen = new IntOpenHashSet();
		while(iit.hasNext())
		{
			int tag = iit.nextInt();
			int num = Integer.parseInt(StringToInt.reverse(tag))/2;
			if(seen.contains(num))
			{continue;}
			seen.add(num);
			int one = StringToInt.ensure(Integer.toString(num*2));
			int two = StringToInt.ensure(Integer.toString(num*2+1));
			int unit;
			this.headToHeadLeftSecUnits.put(one, unit = this.getNextFree());
			this.headToHeadLeftSecUnits.put(two, unit);
			this.headToHeadRightSecUnits.put(one, unit = this.getNextFree());
			this.headToHeadRightSecUnits.put(two, unit);
			this.stopAdjacentLeftSecUnits.put(one, unit = this.getNextFree());
			this.stopAdjacentLeftSecUnits.put(two, unit);
			this.stopAdjacentRightSecUnits.put(one, unit = this.getNextFree());
			this.stopAdjacentRightSecUnits.put(two, unit);
			this.stopDistantLeftSecUnits.put(one, unit = this.getNextFree());
			this.stopDistantLeftSecUnits.put(two, unit);
			this.stopDistantRightSecUnits.put(one, unit = this.getNextFree());
			this.stopDistantRightSecUnits.put(two, unit);
		}
		this.headToHeadLeftSecUnits.put(-1, this.getNextFree());
		this.headToHeadRightSecUnits.put(-1, getNextFree());
		this.stopAdjacentLeftSecUnits.put(-1, getNextFree());
		this.stopAdjacentRightSecUnits.put(-1, getNextFree());
		this.stopDistantLeftSecUnits.put(-1, getNextFree());
		this.stopDistantRightSecUnits.put(-1, getNextFree());
		this.hHSecUnits = new IntOpenHashSet(this.headToHeadLeftSecUnits.values());
		this.hHSecUnits.addAll(this.headToHeadRightSecUnits.values());
		this.sSecUnits = new IntOpenHashSet(this.stopAdjacentLeftSecUnits.values());
		this.sSecUnits.addAll(this.stopAdjacentRightSecUnits.values());
		this.sSecUnits.addAll(this.stopDistantLeftSecUnits.values());
		this.sSecUnits.addAll(this.stopDistantRightSecUnits.values());
	}
	/**
	 * 
	 */
	private final Int2IntOpenHashMap headToHeadLeftSecUnits;
	/**
	 * 
	 * @param head
	 * @param left
	 * @return
	 */
	protected int getHeadToHeadUnit(int head,boolean left)
	{
		if(left)
		{return headToHeadLeftSecUnits.get(head);}
		else
		{return headToHeadRightSecUnits.get(head);}
	}
	/**
	 * 
	 */
	private final Int2IntOpenHashMap headToHeadRightSecUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopAdjacentLeftSecUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopDistantLeftSecUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopAdjacentRightSecUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopDistantRightSecUnits;
	/**
	 * 
	 * @param head
	 * @param left
	 * @param adjacent
	 * @return
	 */
	protected int getSecStopUnit(int head, boolean left, boolean adjacent)
	{
		if(left)
		{
			if(adjacent)
			{return stopAdjacentLeftSecUnits.get(head);}
			else
			{return stopDistantLeftSecUnits.get(head);}
		}
		else
		{
			if(adjacent)
			{return stopAdjacentRightSecUnits.get(head);}
			else
			{return stopDistantRightSecUnits.get(head);}
		}
	}
	/**
	 * 
	 */
	private final IntSet sSecUnits;
	/**
	 * 
	 */
	private final IntSet hHSecUnits;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash hSecUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash sSecUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash hSecElement;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash sSecElement;
	/**
	 * 
	 * @param pss
	 * @param oue
	 * @param thread
	 * @param dir
	 */
	protected void addFeatures(PopulationSamplingStructure pss,
			OpenUnitElementCounter oue, int thread, double dir)
	{
		makeLeftAndRightChildren(pss,thread);
		HeadAssignmentStructure has = (HeadAssignmentStructure) pss;
		WordAndTagHolder wth = (WordAndTagHolder) pss;
		for(int pos=0;pos<has.getLength();++pos)
		{
			int tag = wth.getTag(pos);
			int head = has.getHeadPosition(pos, thread);
			int headTag;
			if(head >= 0)
			{
				if(head == has.getDependencyRootPosition())
				{headTag = -1;}
				else
				{headTag = wth.getTag(head);}
				oue.add(this.getSecHeadToHeadUnit(headTag, pos < head), tag, dir);
			}
			int leftChildren = this.getChildren(true, thread).getInt(pos);
			int rightChildren = this.getChildren(false, thread).getInt(pos);
			makeStops(oue, dir, tag, leftChildren, true);
			makeStops(oue, dir, tag, rightChildren, false);
		}
		super.addFeatures(pss, oue, thread, dir);
	}
	/**
	 * 
	 * @param headTag
	 * @param b
	 * @return
	 */
	private int getSecHeadToHeadUnit(int headTag, boolean left)
	{
		if(left)
		{return this.headToHeadLeftSecUnits.get(headTag);}
		else
		{return this.headToHeadRightSecUnits.get(headTag);}
	}
	/**
	 * @param oue
	 * @param dir
	 * @param tag
	 * @param leftChildren
	 * @param left
	 */
	protected void makeStops(OpenUnitElementCounter oue, double dir, int tag,
			int leftChildren, boolean left)
	{
		if(leftChildren == 0)
		{oue.add(this.getSecStopUnit(tag, left, true), 0, dir);}
		else
		{
			oue.add(this.getSecStopUnit(tag, left, true), 1, dir);
			if(leftChildren > 1)
			{oue.add(this.getSecStopUnit(tag, left, false), 1, dir*(leftChildren-1));}
			oue.add(this.getSecStopUnit(tag, left, false), 0, dir);
		}
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getUnitCash(int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getUnitCash(int unit)
	{
		if(this.hHSecUnits.contains(unit))
		{return this.hSecUnit;}
		else if (this.sSecUnits.contains(unit)) 
		{return this.sSecUnit;}
		return super.getUnitCash(unit);
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getElementCash(int, int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getElementCash(int unit,
			int element)
	{
		if(this.hHSecUnits.contains(unit))
		{return this.hSecElement;}
		else if (this.sSecUnits.contains(unit)) 
		{return this.sSecElement;}
		return super.getElementCash(unit, element);
	}
	@Override
	public double makeChartProposalProb(
			NonProjDepdTree tree,
			int headCandidate, int child, int thread)
	{
		OpenUnitElementCounter oue = getSeen(thread);
		int tag = tree.getTag(child);
		int headTag;
		if(headCandidate == tree.getDependencyRootPosition())
		{headTag = -1;}
		else
		{headTag = tree.getTag(headCandidate);}
		oue.add(this.getSecHeadToHeadUnit(headTag, child < headCandidate), tag, 1.0);
		double d = this.makeLogProb(oue, thread);
		return d + super.makeChartProposalProb(tree, headCandidate, child, thread);
	}
}